// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/dwio/common/options.h>
#include <pollux/plan/plan_builder.h>

namespace kumo::pollux::exec::test {
    /// Contains the query plan and input data files keyed on source plan node ID.
    /// All data files use the same file format specified in 'dataFileFormat'.
    struct TpchPlan {
        core::PlanNodePtr plan;
        std::unordered_map<core::PlanNodeId, std::vector<std::string> > dataFiles;
        dwio::common::FileFormat dataFileFormat;
    };

    /// Contains type information, data files, and file column names for a table.
    /// This information is inferred from the input data files.
    /// The type names are mapped to the standard names.
    /// Example: If the file has a 'returnflag' column, the corresponding type name
    /// will be 'l_returnflag'. fileColumnNames store the mapping between standard
    /// names and the corresponding name in the file.
    struct TpchTableMetadata {
        RowTypePtr type;
        std::vector<std::string> dataFiles;
        std::unordered_map<std::string, std::string> fileColumnNames;
    };

    /// Builds TPC-H queries using TPC-H data files located in the specified
    /// directory. Each table data must be placed in hive-style partitioning. That
    /// is, the top-level directory is expected to contain a sub-directory per table
    /// name and the name of the sub-directory must match the table name. Example:
    /// ls -R data/
    ///  customer   lineitem
    ///
    ///  data/customer:
    ///  customer1.parquet  customer2.parquet
    ///
    ///  data/lineitem:
    ///  lineitem1.parquet  lineitem2.parquet  lineitem3.parquet

    /// The column names can vary. Additional columns may exist towards the end.
    /// The class uses standard names (example: l_returnflag) to build TPC-H plans.
    /// Since the column names in the file can vary, they are mapped to the standard
    /// names. Therefore, the order of the columns in the file is important and
    /// should be in the same order as in the TPC-H standard.
    class TpchQueryBuilder {
    public:
        explicit TpchQueryBuilder(dwio::common::FileFormat format)
            : format_(format) {
        }

        /// Read each data file, initialize row types, and determine data paths for
        /// each table.
        /// @param dataPath path to the data files
        void initialize(const std::string &dataPath);

        /// Get the query plan for a given TPC-H query number.
        /// @param queryId TPC-H query number
        TpchPlan getQueryPlan(int queryId) const;

        /// Returns a plan for select max(c1), max(c2), .. from lineitem where partkey
        /// between
        /// ... such that 'columnPct' columns are aggregated with max(c) for numbers
        /// and with max(length(c)) for strings. To select the columns, we sort them
        /// by column name and take the first columnPct percent.
        TpchPlan getIoMeterPlan(int columnPct) const;

        /// Get the TPC-H table names present.
        static const std::vector<std::string> &getTableNames();

    private:
        // Initializes the schema information for 'tableName' from sample file at
        // 'filePath'.
        void readFileSchema(
            const std::string &tableName,
            const std::string &filePath,
            const std::vector<std::string> &columns);

        TpchPlan getQ1Plan() const;

        TpchPlan getQ2Plan() const;

        TpchPlan getQ3Plan() const;

        TpchPlan getQ4Plan() const;

        TpchPlan getQ5Plan() const;

        TpchPlan getQ6Plan() const;

        TpchPlan getQ7Plan() const;

        TpchPlan getQ8Plan() const;

        TpchPlan getQ9Plan() const;

        TpchPlan getQ10Plan() const;

        TpchPlan getQ11Plan() const;

        TpchPlan getQ12Plan() const;

        TpchPlan getQ13Plan() const;

        TpchPlan getQ14Plan() const;

        TpchPlan getQ15Plan() const;

        TpchPlan getQ16Plan() const;

        TpchPlan getQ17Plan() const;

        TpchPlan getQ18Plan() const;

        TpchPlan getQ19Plan() const;

        TpchPlan getQ20Plan() const;

        TpchPlan getQ21Plan() const;

        TpchPlan getQ22Plan() const;

        const std::vector<std::string> &getTableFilePaths(
            const std::string &tableName) const {
            return tableMetadata_.at(tableName).dataFiles;
        }

        std::shared_ptr<const RowType> getRowType(
            const std::string &tableName,
            const std::vector<std::string> &columnNames) const {
            auto columnSelector = std::make_shared<dwio::common::ColumnSelector>(
                tableMetadata_.at(tableName).type, columnNames);
            return columnSelector->buildSelectedReordered();
        }

        const std::unordered_map<std::string, std::string> &getFileColumnNames(
            const std::string &tableName) const {
            return tableMetadata_.at(tableName).fileColumnNames;
        }

        std::unordered_map<std::string, TpchTableMetadata> tableMetadata_;
        const dwio::common::FileFormat format_;
        static const std::unordered_map<std::string, std::vector<std::string> >
        kTables_;
        static const std::vector<std::string> kTableNames_;

        static constexpr const char *kLineitem = "lineitem";
        static constexpr const char *kCustomer = "customer";
        static constexpr const char *kOrders = "orders";
        static constexpr const char *kNation = "nation";
        static constexpr const char *kRegion = "region";
        static constexpr const char *kPart = "part";
        static constexpr const char *kSupplier = "supplier";
        static constexpr const char *kPartsupp = "partsupp";
        std::shared_ptr<memory::MemoryPool> pool_ =
                memory::memoryManager()->addLeafPool();
    };
} // namespace kumo::pollux::exec::test
